#!/usr/bin/env python3
# coding: utf-8
# File: MedicalGraph.py
# Author: lhy<lhy_in_blcu@126.com,https://huangyong.github.io>
# Date: 18-10-3

import os
import json
from py2neo import Graph,Node,NodeMatcher,Relationship
# 用于计时，提高效率
import time 

from utils.printMessage import (
    logMessage
)

class MedicalGraph:
    def __init__(self):
        # 初始化数据的路径
        cur_dir = os.path.abspath(os.path.dirname(__file__))
        self.data_path = os.path.join(cur_dir, 'data', 'medical.json')
        logMessage("初始化路径")
        # 初始化连接
        self.g = Graph(
            "127.0.0.1:7474",  # neo4j 搭载服务器的ip地址，ifconfig可获取到
            auth=("neo4j", "88888888"),
            name="neo4j")
        # 匹配节点的匹配器
        self.node_matcher = NodeMatcher(self.g)
        logMessage("数据库连接成功")

    def clearGraph(self):
        # 执行Cypher查询删除所有节点和关系
        query = "MATCH (n) DETACH DELETE n"
        self.g.run(query)

    def clearRelationships(self):
        # 执行Cypher查询删除关系,但是会保留节点
        query = "MATCH ()-[r]-() DELETE r"
        self.g.run(query)

    '''读取文件'''
    def read_nodes(self):
        # 共７类节点
        drugs = [] # 药品
        foods = [] #　食物
        checks = [] # 检查
        departments = [] #科室
        producers = [] #药品大类
        diseases = [] #疾病
        symptoms = []#症状

        disease_infos = []#疾病信息

        # 构建节点实体关系
        rels_department = [] #　科室－科室关系
        rels_noteat = [] # 疾病－忌吃食物关系
        rels_doeat = [] # 疾病－宜吃食物关系
        rels_recommandeat = [] # 疾病－推荐吃食物关系
        rels_commonddrug = [] # 疾病－通用药品关系
        rels_recommanddrug = [] # 疾病－热门药品关系
        rels_check = [] # 疾病－检查关系
        rels_drug_producer = [] # 厂商－药物关系

        rels_symptom = [] #疾病症状关系
        rels_acompany = [] # 疾病并发关系
        rels_category = [] #　疾病与科室之间的关系


        count = 0
        for data in open(self.data_path, encoding='utf-8'):
            disease_dict = {}
            count += 1
            if count % 100 == 0:
                logMessage(f"读取第{count}条数据")
            # 转换json数据
            data_json = json.loads(data)
            # 疾病类型存储起来
            disease = data_json['name']
            diseases.append(disease)

            # 初始化各个值
            disease_dict['name'] = disease
            disease_dict['desc'] = ''
            disease_dict['prevent'] = ''
            disease_dict['cause'] = ''
            disease_dict['easy_get'] = ''
            disease_dict['cure_department'] = ''
            disease_dict['cure_way'] = ''
            disease_dict['cure_lasttime'] = ''
            disease_dict['symptom'] = ''
            disease_dict['cured_prob'] = ''

            # 症状
            if 'symptom' in data_json:
                # 数组叠加，存放所有的症状，最后会用set进行过滤
                symptoms += data_json['symptom']
                for symptom in data_json['symptom']:
                    # 用于存储关系边，后续构建
                    rels_symptom.append([disease, symptom])

            if 'acompany' in data_json:
                for acompany in data_json['acompany']:
                    rels_acompany.append([disease, acompany])

            if 'desc' in data_json:
                disease_dict['desc'] = data_json['desc']

            if 'prevent' in data_json:
                disease_dict['prevent'] = data_json['prevent']

            if 'cause' in data_json:
                disease_dict['cause'] = data_json['cause']

            if 'get_prob' in data_json:
                disease_dict['get_prob'] = data_json['get_prob']

            if 'easy_get' in data_json:
                disease_dict['easy_get'] = data_json['easy_get']

            # 治疗部门
            if 'cure_department' in data_json:
                cure_department = data_json['cure_department']
                # 治疗部门最多二级
                # 如果只有一级部门，直接添加疾病和部门的关系
                if len(cure_department) == 1:
                     rels_category.append([disease, cure_department[0]])
                # 如果有二级部门，
                # 不仅要添加疾病和二级部门的关系，
                # 还要添加一级部门和二级部门的关系，上级与下级的关系
                if len(cure_department) == 2:
                    big = cure_department[0]
                    small = cure_department[1]
                    rels_department.append([small, big])
                    rels_category.append([disease, small])

                # 赋值
                disease_dict['cure_department'] = cure_department
                # 部门累加
                departments += cure_department

            if 'cure_way' in data_json:
                disease_dict['cure_way'] = data_json['cure_way']

            if  'cure_lasttime' in data_json:
                disease_dict['cure_lasttime'] = data_json['cure_lasttime']

            if 'cured_prob' in data_json:
                disease_dict['cured_prob'] = data_json['cured_prob']

            # 可治疗的药物
            if 'common_drug' in data_json:
                common_drug = data_json['common_drug']
                for drug in common_drug:
                    rels_commonddrug.append([disease, drug])
                drugs += common_drug

            if 'recommand_drug' in data_json:
                recommand_drug = data_json['recommand_drug']
                drugs += recommand_drug
                for drug in recommand_drug:
                    rels_recommanddrug.append([disease, drug])

            # not_eat的话表示有食物，说明do_eat、recommand_eat都有
            if 'not_eat' in data_json:
                not_eat = data_json['not_eat']
                for _not in not_eat:
                    rels_noteat.append([disease, _not])

                foods += not_eat
                do_eat = data_json['do_eat']
                for _do in do_eat:
                    rels_doeat.append([disease, _do])

                foods += do_eat
                recommand_eat = data_json['recommand_eat']

                for _recommand in recommand_eat:
                    rels_recommandeat.append([disease, _recommand])
                foods += recommand_eat

            if 'check' in data_json:
                check = data_json['check']
                for _check in check:
                    rels_check.append([disease, _check])
                checks += check
            if 'drug_detail' in data_json:
                drug_detail = data_json['drug_detail']
                # drug_detail=["欧立停(苯甲酸利扎曲普坦片)", ......]
                # 获取生产商
                producer = [i.split('(')[0] for i in drug_detail]
                # 构建和药物生产商之间的关系
                rels_drug_producer += [[i.split('(')[0], i.split('(')[-1].replace(')', '')] for i in drug_detail]
                producers += producer
            disease_infos.append(disease_dict)
        return set(drugs), set(foods), set(checks), set(departments), set(producers), set(symptoms), set(diseases), disease_infos,\
               rels_check, rels_recommandeat, rels_noteat, rels_doeat, rels_department, rels_commonddrug, rels_drug_producer, rels_recommanddrug,\
               rels_symptom, rels_acompany, rels_category

    '''建立节点'''
    def create_node(self, label, nodes, nodeTypeName):
        # 提高效率采用130一组进行提交
        summitCount = 130
        count = 0
        total = len(nodes)
        tx = self.g.begin()
        for node_name in nodes:
            node = Node(label, name=node_name)
            tx.create(node)
            count += 1
            if count % 100 == 0:
                logMessage(f"创建{nodeTypeName}节点,完成{count}个,共{total}")
            # 到达一定的数量就提交
            if count % summitCount == 0:
                self.g.commit(tx)
                tx = self.g.begin()
        self.g.commit(tx)
        logMessage(f"创建知识图谱中心{nodeTypeName}的节点成功")

    '''创建知识图谱中心疾病的节点'''
    def create_diseases_nodes(self, disease_infos) -> None:
        # 提高效率采用130一组进行提交
        summitCount = 130
        count = 0
        total = len(disease_infos)
        tx = self.g.begin()
        for disease_dict in disease_infos:
            node = Node("Disease", name=disease_dict['name'], desc=disease_dict['desc'],
                        prevent=disease_dict['prevent'] ,cause=disease_dict['cause'],
                        easy_get=disease_dict['easy_get'],cure_lasttime=disease_dict['cure_lasttime'],
                        cure_department=disease_dict['cure_department']
                        ,cure_way=disease_dict['cure_way'] , cured_prob=disease_dict['cured_prob'])
            tx.create(node)
            count += 1
            if count % 100 == 0:
                logMessage(f"创建知识图谱中心疾病的节点,完成{count}个,共{total}个")
            # 到达一定的数量就提交
            if count % summitCount == 0:
                self.g.commit(tx)
                tx = self.g.begin()  
        self.g.commit(tx)
        logMessage("创建知识图谱中心疾病的节点成功")
        return

    '''创建知识图谱实体节点类型schema'''
    def create_graphnodes(self):
        # 获取数组,用于构建节点和实体关系
        Drugs, Foods, Checks, Departments, Producers, \
        Symptoms, Diseases, disease_infos, rels_check, \
        rels_recommandeat, rels_noteat, rels_doeat, rels_department, \
        rels_commonddrug, rels_drug_producer, rels_recommanddrug,\
        rels_symptom, rels_acompany, rels_category = self.read_nodes()
        # print("原始项目")
        # print(f"疾病节点数量:{len(disease_infos)}")
        # print(f"食物节点数量:{len(Foods)}")
        # print(f"药物节点数量:{len(Drugs)}")
        # print(f"症状节点数量:{len(Symptoms)}")
        # print(f"检查节点数量:{len(Checks)}")
        # print(f"生产商节点数量:{len(Producers)}")
        # print(f"诊断部门节点数量:{len(Departments)}")
        # return
        testList = []
        # 1.创建疾病节点
        self.create_diseases_nodes(disease_infos)
        # 2.创建药物节点
        self.create_node('Drug', Drugs, nodeTypeName="药物")
        # 3.创建食物节点
        self.create_node('Food', Foods, nodeTypeName="食物")
        # 4.创建检查节点
        self.create_node('Check', Checks, nodeTypeName="检查")
        # 5.创建部门节点
        self.create_node('Department', Departments, nodeTypeName="部门")
        # 6.创建生产商节点
        self.create_node('Producer', Producers, nodeTypeName="生产商")
        # 7.创建症状商节点
        self.create_node('Symptom', Symptoms, nodeTypeName="症状")
        return


    '''创建实体关系边'''
    def create_graphrels(self):
        Drugs, Foods, Checks, Departments, Producers, Symptoms, Diseases, disease_infos, rels_check, rels_recommandeat, rels_noteat, rels_doeat, rels_department, rels_commonddrug, rels_drug_producer, rels_recommanddrug,rels_symptom, rels_acompany, rels_category = self.read_nodes()
        # 这个顺序 291112条边
        # 耗时2878秒
        # self.create_relationship('Disease', 'Food', rels_recommandeat, 'recommand_eat', '推荐食谱')
        # self.create_relationship('Disease', 'Food', rels_noteat, 'no_eat', '忌吃')
        # self.create_relationship('Disease', 'Food', rels_doeat, 'do_eat', '宜吃')
        # self.create_relationship('Disease', 'Symptom', rels_symptom, 'has_symptom', '症状')
        # self.create_relationship('Disease', 'Disease', rels_acompany, 'acompany_with', '并发症')
        # self.create_relationship('Department', 'Department', rels_department, 'belongs_to', '属于')
        # self.create_relationship('Disease', 'Drug', rels_commonddrug, 'common_drug', '常用药品')
        # self.create_relationship('Disease', 'Check', rels_check, 'need_check', '诊断检查')
        # self.create_relationship('Producer', 'Drug', rels_drug_producer, 'drugs_of', '生产药品')
        # self.create_relationship('Disease', 'Drug', rels_recommanddrug, 'recommand_drug', '好评药品')
        # self.create_relationship('Disease', 'Department', rels_category, 'belongs_to', '所属科室')

        # 这个顺序
        # 耗时3089
        self.create_relationship('Disease', 'Food', rels_recommandeat, 'recommand_eat', '推荐食谱')
        self.create_relationship('Disease', 'Food', rels_noteat, 'no_eat', '忌吃')
        self.create_relationship('Disease', 'Food', rels_doeat, 'do_eat', '宜吃')
        self.create_relationship('Department', 'Department', rels_department, 'belongs_to', '属于')
        self.create_relationship('Disease', 'Drug', rels_commonddrug, 'common_drug', '常用药品')
        self.create_relationship('Producer', 'Drug', rels_drug_producer, 'drugs_of', '生产药品')
        self.create_relationship('Disease', 'Drug', rels_recommanddrug, 'recommand_drug', '好评药品')
        self.create_relationship('Disease', 'Check', rels_check, 'need_check', '诊断检查')
        self.create_relationship('Disease', 'Symptom', rels_symptom, 'has_symptom', '症状')
        self.create_relationship('Disease', 'Disease', rels_acompany, 'acompany_with', '并发症')
        self.create_relationship('Disease', 'Department', rels_category, 'belongs_to', '所属科室')

    '''创建实体关联边'''
    def create_relationship(self, start_node_type, end_node_type, edges, rel_type, rel_name):
        summitCount = 300
        count = 0
        tx = self.g.begin()
        # 去重处理,这里使用到了set，
        # set的类型要求整数、浮点数、字符串、元组等
        # 列表不支持，因此需要把[a, b]变成 a###b
        set_edges = []
        for edge in edges:
            set_edges.append('###'.join(edge))
        total = len(set(set_edges))
        for edge in set(set_edges):
            count += 1
            edge = edge.split('###')
            start_node_name = edge[0]
            end_node_name = edge[1]
            # 先将节点匹配出来
            start_node = self.node_matcher.match(start_node_type, name=start_node_name).first()
            end_node = self.node_matcher.match(end_node_type, name=end_node_name).first()
            # 使用 CREATE 语句创建关系
            relation = Relationship(start_node, rel_type, end_node, name=rel_name)
            tx.create(relation)
            # query = "match(p:%s),(q:%s) where p.name='%s'and q.name='%s' create (p)-[rel:%s{name:'%s'}]->(q)" % (
            #     start_node, end_node, p, q, rel_type, rel_name)
            # tx.run(query)
            if count % 100 == 0:
                logMessage(f"创建知识图谱中心疾病的实体关系{rel_name},完成{count}个,共{total}个")
            if count % summitCount == 0:
                self.g.commit(tx)
                tx = self.g.begin()
        self.g.commit(tx)
        logMessage(f"创建知识图谱中心疾病的实体关系{rel_name}成功")
        return

    '''导出数据'''
    def export_data(self):
        Drugs, Foods, Checks, Departments, Producers, Symptoms, Diseases, disease_infos, rels_check, rels_recommandeat, rels_noteat, rels_doeat, rels_department, rels_commonddrug, rels_drug_producer, rels_recommanddrug, rels_symptom, rels_acompany, rels_category = self.read_nodes()
        f_drug = open('drug.txt', 'w+')
        f_food = open('food.txt', 'w+')
        f_check = open('check.txt', 'w+')
        f_department = open('department.txt', 'w+')
        f_producer = open('producer.txt', 'w+')
        f_symptom = open('symptoms.txt', 'w+')
        f_disease = open('disease.txt', 'w+')

        f_drug.write('\n'.join(list(Drugs)))
        f_food.write('\n'.join(list(Foods)))
        f_check.write('\n'.join(list(Checks)))
        f_department.write('\n'.join(list(Departments)))
        f_producer.write('\n'.join(list(Producers)))
        f_symptom.write('\n'.join(list(Symptoms)))
        f_disease.write('\n'.join(list(Diseases)))

        f_drug.close()
        f_food.close()
        f_check.close()
        f_department.close()
        f_producer.close()
        f_symptom.close()
        f_disease.close()

        return



if __name__ == '__main__':
    start_time = time.time()
    handler = MedicalGraph()
    logMessage("step1:清空图谱")
    handler.clearGraph()
    logMessage("step2:导入图谱节点中")
    handler.create_graphnodes()
    logMessage("step3:导入图谱边中")   
    # handler.create_graphrels()
    end_time = time.time()
    spend_time = end_time - start_time
    logMessage(f"花费的总时间为{spend_time}")
    
